#!/usr/bin/python3
# coding=utf-8
#
# Copyright (C) 2025. Huawei Technologies Co., Ltd. All rights reserved.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# ===============================================================================
import torch
from xlite._C import Runtime, add


rt = Runtime(0, 500)
torch.npu.set_device(0)

supported_dtype_list = [torch.float16, torch.bfloat16]

for dtype in supported_dtype_list:
    x = torch.randn(8, 2048, dtype=dtype, device="npu:0")
    y = torch.randn(8, 2048, dtype=dtype, device="npu:0")
    z = torch.empty(8, 2048, dtype=dtype, device="npu:0")

    standard = x + y

    torch.npu.synchronize()
    add(rt, x, y, z)
    torch.npu.synchronize()
    print(f'add {dtype} executed!')

    try:
        torch.testing.assert_close(standard, z, atol=1e-5, rtol=1e-3)
    except AssertionError as e:
        print(f'{e}')
        print(f'x: {x}')
        print(f'y: {y}')
        print(f'torch_npu: {standard}')
        print(f'xlite: {z}')